use core::{
    any::Any,
    hint::spin_loop,
    sync::atomic::{AtomicBool, Ordering},
};

use axerrno::LinuxError;
use axfs_ng_vfs::{NodeFlags, VfsResult};
use axhal::mem::phys_to_virt;
use kspin::SpinNoIrq;
use sbox_proto::timer::{self, ChannelRequest, MicrosRequest, WaitRequest};
use spin::Once;
use starry_vm::{VmMutPtr, VmPtr};
use timer_rk3588::{RK3588Timer, TimerMode};

use crate::vfs::DeviceOps;

const TIMER_FREQ_HZ: u64 = 24_000_000;

struct TimerChannel {
    timer: SpinNoIrq<RK3588Timer>,
    initialised: AtomicBool,
}

impl TimerChannel {
    fn new(base: usize, channel: u8) -> Self {
        Self {
            timer: SpinNoIrq::new(RK3588Timer::new(base, channel)),
            initialised: AtomicBool::new(false),
        }
    }

    fn with_timer<T>(&self, f: impl FnOnce(&mut RK3588Timer) -> VfsResult<T>) -> VfsResult<T> {
        let mut guard = self.timer.lock();
        if !self.initialised.load(Ordering::Acquire) {
            if !guard.init() {
                return Err(LinuxError::ENODEV);
            }
            self.initialised.store(true, Ordering::Release);
        }
        f(&mut guard)
    }

    fn disable(&self) -> VfsResult<()> {
        self.with_timer(|timer| {
            timer.disable();
            Ok(())
        })
    }

    fn wait_ticks(&self, ticks: u64) -> VfsResult<u64> {
        self.with_timer(|timer| {
            timer.disable();
            timer.clear_interrupt();

            if ticks == 0 {
                return Ok(0);
            }

            timer.load_count(ticks);
            timer.set_mode(TimerMode::UserDefined);
            timer.disable_interrupt();
            timer.enable();

            let mut last_value = timer.current_value();
            let mut stagnant_cycles = 0usize;

            loop {
                if timer.get_interrupt_status() || !timer.is_enabled() {
                    timer.disable();
                    timer.clear_interrupt();
                    return Ok(ticks);
                }

                let current = timer.current_value();

                if current != last_value {
                    last_value = current;
                    stagnant_cycles = 0;
                } else {
                    stagnant_cycles = stagnant_cycles.saturating_add(1);
                }

                if stagnant_cycles > 1_000_000 {
                    timer.disable();
                    timer.clear_interrupt();
                    return Err(LinuxError::ETIMEDOUT);
                }

                spin_loop();
            }
        })
    }

    fn wait_micros(&self, micros: u64) -> VfsResult<u64> {
        let ticks = micros
            .checked_mul(TIMER_FREQ_HZ)
            .ok_or(LinuxError::EINVAL)?
            / 1_000_000;
        self.wait_ticks(ticks)
    }
}

struct TimerChannels {
    channels: [TimerChannel; 2],
}

impl TimerChannels {
    fn new() -> Self {
        use timer_rk3588::regs::base_addr::TIMER_PMU;

        let base = phys_to_virt(TIMER_PMU.into()).as_usize();

        Self {
            channels: [TimerChannel::new(base, 0), TimerChannel::new(base, 1)],
        }
    }

    fn get(&self, index: usize) -> VfsResult<&TimerChannel> {
        self.channels.get(index).ok_or(LinuxError::EINVAL)
    }
}

static TIMER_CHANNELS: Once<TimerChannels> = Once::new();

fn channels() -> &'static TimerChannels {
    TIMER_CHANNELS.call_once(TimerChannels::new)
}

fn get_channel(channel: u32) -> VfsResult<&'static TimerChannel> {
    channels().get(channel as usize)
}

/// Character device that exposes the RK3588 PMU timer to user space.
pub struct TimerDevice;

impl TimerDevice {
    pub const fn new() -> Self {
        Self
    }

    fn handle_init(request: &ChannelRequest) -> VfsResult<()> {
        get_channel(request.channel)?.with_timer(|_| Ok(()))
    }

    fn handle_disable(request: &ChannelRequest) -> VfsResult<()> {
        get_channel(request.channel)?.disable()
    }

    fn handle_wait_ticks(ptr: *mut WaitRequest) -> VfsResult<()> {
        if ptr.is_null() {
            return Err(LinuxError::EINVAL);
        }
        let mut req = unsafe { ptr.vm_read_uninit()?.assume_init() };
        let observed = get_channel(req.channel)?.wait_ticks(req.ticks)?;
        req.observed = observed;
        ptr.vm_write(req)?;
        Ok(())
    }

    fn handle_wait_micros(ptr: *mut MicrosRequest) -> VfsResult<()> {
        if ptr.is_null() {
            return Err(LinuxError::EINVAL);
        }
        let mut req = unsafe { ptr.vm_read_uninit()?.assume_init() };
        let observed = get_channel(req.channel)?.wait_micros(req.micros)?;
        req.observed = observed;
        ptr.vm_write(req)?;
        Ok(())
    }
}

impl Default for TimerDevice {
    fn default() -> Self {
        Self::new()
    }
}

impl DeviceOps for TimerDevice {
    fn read_at(&self, _buf: &mut [u8], _offset: u64) -> VfsResult<usize> {
        Ok(0)
    }

    fn write_at(&self, buf: &[u8], _offset: u64) -> VfsResult<usize> {
        Ok(buf.len())
    }

    fn ioctl(&self, cmd: u32, arg: usize) -> VfsResult<usize> {
        match cmd {
            timer::IOCTL_INIT => {
                let req = unsafe {
                    (arg as *const ChannelRequest)
                        .vm_read_uninit()?
                        .assume_init()
                };
                Self::handle_init(&req)?;
            }
            timer::IOCTL_DISABLE => {
                let req = unsafe {
                    (arg as *const ChannelRequest)
                        .vm_read_uninit()?
                        .assume_init()
                };
                Self::handle_disable(&req)?;
            }
            timer::IOCTL_WAIT_TICKS => {
                Self::handle_wait_ticks(arg as *mut WaitRequest)?;
            }
            timer::IOCTL_WAIT_MICROS => {
                Self::handle_wait_micros(arg as *mut MicrosRequest)?;
            }
            _ => return Err(LinuxError::ENOTTY),
        }
        Ok(0)
    }

    fn as_any(&self) -> &dyn Any {
        self
    }

    fn flags(&self) -> NodeFlags {
        NodeFlags::NON_CACHEABLE | NodeFlags::STREAM
    }
}
